import av
import numpy as np
import cv2

from demo.demo import Predictor
from nanodet.util import load_config, cfg, Logger, overlay_bbox_cv


class videoTansNanodet:

    def __init__(self,config_path,model_path,device='cpu'):

        #1.load config
        load_config(cfg, config_path)
        self.cfg=cfg
        logger = Logger(-1, use_tensorboard=False)

        self.predictor = Predictor(cfg, model_path, logger, device=device)

    def transform(self, frame: av.VideoFrame) -> np.ndarray:
        img = frame.to_ndarray(format="bgr24")
        meta, results = self.predictor.inferenceWithLoadImg(img)
        return overlay_bbox_cv(img, results, self.cfg.class_names, 0.35)



# if __name__=="__main__":
#
#     config_path = '../nanodet/config/nanodet-m.yml'
#     model_path = '../nanodet/nanodet_m.ckpt'
#     videoNanodet=videoTansNanodet(config_path,model_path)
#     img=cv2.imread("../data/img/share.jpg")
#     videoNanodet.getNanodetDetect(img)
